// SPDX-License-Identifier: Mulan PSL v2
/*
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This software is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *         http://license.coscl.org.cn/MulanPSL2
 *
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

use std::{
    env,
    ffi::OsStr,
    fmt::Debug,
    fs,
    os::unix::{ffi::OsStrExt, fs::symlink},
    path::{Path, PathBuf},
    process::{Command, Stdio},
};

use anyhow::{Context, Result};
use indexmap::IndexSet;
use once_cell::sync::Lazy;
use walkdir::WalkDir;

const CUDA_ENV_NAME: &str = "CUDA_HOME";
const NCCL_ENV_NAME: &str = "NCCL_HOME";
const VIRT_ENV_NAME: &str = "VIRTUAL_ENV";

const DEFAULT_CUDA_HOME: &str = "/usr/local/cuda";
const DEFAULT_NCCL_HOME: &str = "/usr/local/nccl";

const MAX_SEARCH_DEPTH: usize = 5;

const CPP_COMPILER_ARGS: &[&str] = &["-x", "c++", "-std=c++17"];

const STATIC_LIB_SUFFIX: &str = ".a";
const DYNAMIC_LIB_PREFIX: &str = "lib";
const DYNAMIC_LIB_SUFFIX: &str = ".so";

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum SearchOptions {
    MatchOne,
    MatchAll,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum Language {
    C,
    Cxx,
}

impl Language {
    fn is_cpp(&self) -> bool {
        matches!(self, Self::Cxx)
    }
}

#[derive(Debug)]
struct Binding {
    src_name: &'static str,
    lib_name: Option<&'static str>,
    language: Language,
    allowlist: &'static [&'static str],
    output_name: &'static str,
}

#[derive(Debug)]
struct BuildEntry {
    src_file: PathBuf,
    lib_file: Option<PathBuf>,
    binding: Binding,
}

#[derive(Debug)]
struct BuildEnvironment {
    include_dirs: Vec<PathBuf>,
    library_dirs: Vec<PathBuf>,
    output_dir: PathBuf,
}

impl BuildEnvironment {
    fn detect() -> Result<Self> {
        println!("cargo:rerun-if-env-changed={}", CUDA_ENV_NAME);
        println!("cargo:rerun-if-env-changed={}", NCCL_ENV_NAME);
        println!("cargo:rerun-if-env-changed={}", VIRT_ENV_NAME);

        println!("include_dirs:");
        let include_dirs = Self::detect_include_dirs();
        for dir in &include_dirs {
            println!("{}", dir.display());
        }
        println!();

        println!("library_dirs:");
        let library_dirs = Self::detect_library_dirs();
        for dir in &library_dirs {
            println!("{}", dir.display());
        }
        println!();

        println!("output_dir:");
        let output_dir = env::var("OUT_DIR")
            .map(PathBuf::from)
            .context("Failed to get OUT_DIR environment variable")?;
        println!("{}", output_dir.display());
        println!();

        Ok(Self {
            include_dirs,
            library_dirs,
            output_dir,
        })
    }

    fn detect_include_dirs() -> Vec<PathBuf> {
        Self::detect_search_dirs(
            &["include", "src/include", "src/include/plugin"],
            Self::detect_python_nvidia_include_dirs,
            Self::detect_compiler_include_dirs,
            SearchOptions::MatchAll,
        )
    }

    fn detect_library_dirs() -> Vec<PathBuf> {
        Self::detect_search_dirs(
            &["lib64", "lib"],
            Self::detect_python_nvidia_library_dirs,
            Self::detect_compiler_library_dirs,
            SearchOptions::MatchOne,
        )
    }

    fn detect_search_dirs(
        subdirs: &[&str],
        python_dirs_fn: fn() -> Vec<PathBuf>,
        compiler_dirs_fn: fn(Language) -> Vec<PathBuf>,
        options: SearchOptions,
    ) -> Vec<PathBuf> {
        let mut results = IndexSet::new();

        let env_roots: Vec<_> = [CUDA_ENV_NAME, NCCL_ENV_NAME]
            .iter()
            .filter_map(|key| env::var(key).ok())
            .collect();
        Self::find_paths_from_roots(&mut results, &env_roots, subdirs, options);

        results.extend(python_dirs_fn());

        let default_roots = [DEFAULT_CUDA_HOME, DEFAULT_NCCL_HOME];
        Self::find_paths_from_roots(&mut results, &default_roots, subdirs, options);

        results.extend(compiler_dirs_fn(Language::C));
        results.extend(compiler_dirs_fn(Language::Cxx));

        results.into_iter().collect()
    }

    fn find_paths_from_roots<P: AsRef<Path>>(
        results: &mut IndexSet<PathBuf>,
        roots: &[P],
        subdirs: &[&str],
        options: SearchOptions,
    ) {
        for root in roots {
            for subdir in subdirs {
                let path = root.as_ref().join(subdir);
                if path.is_dir() {
                    results.insert(path);
                    if options == SearchOptions::MatchOne {
                        break;
                    }
                }
            }
        }
    }

    fn detect_python_nvidia_dirs() -> (&'static [PathBuf], &'static [PathBuf]) {
        static NVIDIA_DIRS: Lazy<(Vec<PathBuf>, Vec<PathBuf>)> = Lazy::new(|| {
            const PATTERNS: &[&str] = &["nvidia/**/include", "nvidia/**/lib"];
            const PY_SCRIPT: &str = "import sysconfig; print(sysconfig.get_path('platlib'))";

            let result = Command::new("python3").arg("-c").arg(PY_SCRIPT).output();
            let output = match result {
                Ok(out) => out,
                Err(_) => return Default::default(),
            };

            if !output.status.success() {
                return Default::default();
            }

            let package_dirs: Vec<PathBuf> = output
                .stdout
                .split(|&b| b == b'\n')
                .filter(|line| !line.is_empty())
                .map(OsStr::from_bytes)
                .map(PathBuf::from)
                .collect();

            let mut include_dirs = IndexSet::new();
            let mut library_dirs = IndexSet::new();
            for package_dir in package_dirs {
                if !package_dir.is_dir() {
                    continue;
                }

                for pattern in PATTERNS {
                    let full_pattern = package_dir.join(pattern).to_string_lossy().to_string();
                    if let Ok(entries) = glob::glob(&full_pattern) {
                        for entry in entries.filter_map(Result::ok) {
                            // We modified nccl, thus we cannot use nccl from nvidia pip package
                            if entry.to_string_lossy().contains("nccl") {
                                continue;
                            }
                            if entry.is_dir() {
                                if pattern.ends_with("include") {
                                    include_dirs.insert(entry);
                                } else if pattern.ends_with("lib") {
                                    library_dirs.insert(entry);
                                }
                            }
                        }
                    }
                }
            }

            (
                include_dirs.into_iter().collect(),
                library_dirs.into_iter().collect(),
            )
        });

        (NVIDIA_DIRS.0.as_ref(), NVIDIA_DIRS.1.as_ref())
    }

    fn detect_python_nvidia_include_dirs() -> Vec<PathBuf> {
        Self::detect_python_nvidia_dirs().0.to_owned()
    }

    fn detect_python_nvidia_library_dirs() -> Vec<PathBuf> {
        Self::detect_python_nvidia_dirs().1.to_owned()
    }

    fn detect_compiler_include_dirs(lang: Language) -> Vec<PathBuf> {
        let compiler = cc::Build::new().cpp(lang.is_cpp()).get_compiler();
        let lang_flag = if lang.is_cpp() { "-xc++" } else { "-xc" };

        let output = match compiler
            .to_command()
            .env("LANG", "C")
            .args(["-E", lang_flag, "-", "-v"])
            .stdin(Stdio::null())
            .output()
        {
            Ok(out) => out,
            Err(_) => {
                return Default::default();
            }
        };

        output
            .stderr
            .split(|&b| b == b'\n')
            .map(|line| line.trim_ascii())
            .skip_while(|&line| line != b"#include <...> search starts here:")
            .skip(1)
            .take_while(|&line| line != b"End of search list.")
            .map(|slice| PathBuf::from(OsStr::from_bytes(slice)))
            .filter(|p| p.is_dir())
            .collect::<IndexSet<_>>()
            .into_iter()
            .collect()
    }

    fn detect_compiler_library_dirs(lang: Language) -> Vec<PathBuf> {
        let compiler = cc::Build::new().cpp(lang.is_cpp()).get_compiler();

        let output = match compiler.to_command().arg("-print-search-dirs").output() {
            Ok(out) => out,
            Err(_) => {
                return Default::default();
            }
        };

        output
            .stdout
            .split(|&b| b == b'\n')
            .map(|line| line.trim_ascii())
            .find_map(|line_bytes| line_bytes.strip_prefix(b"libraries: ="))
            .into_iter()
            .flat_map(|slice| env::split_paths(OsStr::from_bytes(slice)))
            .filter_map(|path| path.canonicalize().ok())
            .filter(|path| path.is_dir())
            .collect::<IndexSet<_>>()
            .into_iter()
            .collect()
    }
}

impl BuildEntry {
    fn resolve(env: &BuildEnvironment, binding: Binding) -> Result<Self> {
        println!("{}:", binding.src_name);

        let src_file = Self::find_header_file(&binding, &env.include_dirs)
            .with_context(|| format!("Cannot find source file '{}'", binding.src_name))?;
        println!("{}", src_file.display());

        let lib_file = Self::find_library_file(&binding, &env.library_dirs);
        if let Some(lib) = &lib_file {
            println!("{}", lib.display());
        }
        println!();

        let entry = Self {
            src_file,
            lib_file,
            binding,
        };

        Ok(entry)
    }

    fn find_header_file(binding: &Binding, include_dirs: &[PathBuf]) -> Option<PathBuf> {
        let file_name = binding.src_name;
        let match_pattern = format!("**/{}", file_name);

        Self::match_file(&match_pattern, include_dirs, MAX_SEARCH_DEPTH)
    }

    fn find_library_file(binding: &Binding, library_dirs: &[PathBuf]) -> Option<PathBuf> {
        let lib_name = binding.lib_name?;
        let match_patterns = [
            format!("**/lib{}.so*", lib_name),
            format!("**/lib{}.a", lib_name),
        ];

        match_patterns
            .iter()
            .find_map(|pattern| Self::match_file(pattern, library_dirs, MAX_SEARCH_DEPTH))
    }

    fn match_file(pattern: &str, search_dirs: &[PathBuf], max_depth: usize) -> Option<PathBuf> {
        let pattern_glob = glob::Pattern::new(pattern).ok()?;

        search_dirs.iter().find_map(|dir| {
            WalkDir::new(dir)
                .max_depth(max_depth)
                .into_iter()
                .filter_map(|entry| entry.ok())
                .find(|entry| {
                    entry.file_type().is_file()
                        && pattern_glob.matches(&entry.file_name().to_string_lossy())
                })
                .map(|entry| entry.into_path())
        })
    }
}

struct Builder {
    env: BuildEnvironment,
    entries: Vec<BuildEntry>,
}

impl Builder {
    fn new(env: BuildEnvironment, entries: Vec<BuildEntry>) -> Self {
        Self { env, entries }
    }

    fn generate_file(&self, entry: &BuildEntry) -> Result<()> {
        println!("cargo:rerun-if-changed={}", entry.src_file.display());

        let mut builder = bindgen::Builder::default()
            .header(entry.src_file.to_string_lossy())
            .generate_inline_functions(true)
            .generate_comments(false)
            .parse_callbacks(Box::new(bindgen::CargoCallbacks::new()));

        for dir in &self.env.include_dirs {
            builder = builder.clang_arg(format!("-I{}", dir.display()));
        }

        for &pattern in entry.binding.allowlist {
            builder = builder
                .allowlist_function(pattern)
                .allowlist_type(pattern)
                .allowlist_var(pattern);
        }

        if entry.binding.language.is_cpp() {
            builder = builder.clang_args(CPP_COMPILER_ARGS);
        }

        let generated = builder
            .generate()
            .context("Bindgen failed to generate bindings")?;

        let output_file = self.env.output_dir.join(entry.binding.output_name);
        generated
            .write_to_file(&output_file)
            .with_context(|| format!("Failed to write to '{}'", output_file.display()))?;

        Ok(())
    }

    fn generate_files(self) -> Result<Self> {
        fs::create_dir_all(&self.env.output_dir).with_context(|| {
            format!(
                "Failed to create output directory '{}'",
                self.env.output_dir.display()
            )
        })?;

        for entry in &self.entries {
            self.generate_file(entry)
                .with_context(|| format!("{:#?}", self.env))
                .with_context(|| format!("Failed to generate file '{}'", entry.binding.src_name))?;
        }

        Ok(self)
    }

    fn report_link_info(self) -> Result<Self> {
        println!(
            "cargo:rustc-link-search=native={}",
            self.env.output_dir.display()
        );

        for entry in &self.entries {
            if entry.binding.language.is_cpp() {
                println!("cargo:rustc-link-lib=stdc++");
            }

            if let (Some(lib_name), Some(lib_file)) = (entry.binding.lib_name, &entry.lib_file) {
                if let Some(lib_dir) = lib_file.parent() {
                    println!("cargo:rustc-link-search=native={}", lib_dir.display());
                }
                let file_name =
                    lib_file
                        .file_name()
                        .and_then(|s| s.to_str())
                        .with_context(|| {
                            format!(
                                "Library path has a non-UTF8 filename: {}",
                                lib_file.display()
                            )
                        })?;
                if file_name.ends_with(STATIC_LIB_SUFFIX) {
                    println!("cargo:rustc-link-lib=static={}", lib_name);
                } else if file_name.contains(DYNAMIC_LIB_SUFFIX) {
                    let link_name =
                        format!("{}{}{}", DYNAMIC_LIB_PREFIX, lib_name, DYNAMIC_LIB_SUFFIX);
                    let link_path = self.env.output_dir.join(&link_name);

                    if !link_path.exists() {
                        symlink(lib_file, &link_path).with_context(|| {
                            format!(
                                "Failed to create symlink from '{}' to '{}'",
                                lib_file.display(),
                                link_path.display()
                            )
                        })?;
                    }
                    println!("cargo:rustc-link-lib={}", lib_name);
                }
            }
        }

        Ok(self)
    }
}

/// Main entry point for the build script.
fn main() -> Result<()> {
    let bindings = vec![
        #[cfg(feature = "driver")]
        Binding {
            src_name: "cuda.h",
            lib_name: if cfg!(feature = "link") {
                Some("cuda")
            } else {
                None
            },
            language: Language::C,
            allowlist: &["^cu.*"],
            output_name: "driver.rs",
        },
        #[cfg(feature = "runtime")]
        Binding {
            src_name: "cuda_runtime.h",
            lib_name: if cfg!(feature = "link") {
                Some("cudart")
            } else {
                None
            },
            language: Language::C,
            allowlist: &["^cuda.*"],
            output_name: "runtime.rs",
        },
        // ... (rest of the bindings are unchanged)
        #[cfg(feature = "nvml")]
        Binding {
            src_name: "nvml.h",
            lib_name: if cfg!(feature = "link") {
                Some("nvidia-ml")
            } else {
                None
            },
            language: Language::C,
            allowlist: &["^nvml.*"],
            output_name: "nvml.rs",
        },
        #[cfg(feature = "nvrtc")]
        Binding {
            src_name: "nvrtc.h",
            lib_name: if cfg!(feature = "link") {
                Some("nvrtc")
            } else {
                None
            },
            language: Language::C,
            allowlist: &["^nvrtc.*", "^NVRTC_.*"],
            output_name: "nvrtc.rs",
        },
        #[cfg(feature = "cublas")]
        Binding {
            src_name: "cublas.h",
            lib_name: if cfg!(feature = "link") {
                Some("cublas")
            } else {
                None
            },
            language: Language::C,
            allowlist: &["^cublas.*", "^CUBLAS.*"],
            output_name: "cublas.rs",
        },
        #[cfg(feature = "cublaslt")]
        Binding {
            src_name: "cublasLt.h",
            lib_name: if cfg!(feature = "link") {
                Some("cublasLt")
            } else {
                None
            },
            language: Language::C,
            allowlist: &["^cublasLt.*"],
            output_name: "cublaslt.rs",
        },
        #[cfg(feature = "nccl")]
        Binding {
            src_name: "nccl.h",
            lib_name: if cfg!(feature = "link") {
                Some("nccl")
            } else {
                None
            },
            language: Language::C,
            allowlist: &["^nccl.*", "^pnccl.*"],
            output_name: "nccl.rs",
        },
        #[cfg(feature = "bootstrap")]
        Binding {
            src_name: "bootstrap.h",
            lib_name: if cfg!(feature = "link") {
                Some("nccl_static")
            } else {
                None
            },
            language: Language::Cxx,
            allowlist: &["^bootstrap.*", "ncclCalloc.*", "ncclComm", "ncclComm_t"],
            output_name: "bootstrap.rs",
        },
    ];

    println!("------ BUILD_ENVIRONMENT ------");
    let env = BuildEnvironment::detect().context("Failed to detect build environment")?;

    println!("------ BUILD_ENTRY ------");
    let files = bindings
        .into_iter()
        .map(|binding| BuildEntry::resolve(&env, binding))
        .collect::<Result<Vec<_>>>()
        .with_context(|| format!("{:#?}", env))
        .context("Failed to resolve build entries")?;

    Builder::new(env, files)
        .generate_files()?
        .report_link_info()?;

    Ok(())
}
